# coding:utf-8


from deep_env import Maze
from deep_q_learning import DeepQLearning


def update():
    step = 0
    for episode in range(100):
        # 初始化 state(状态)
        state = env.reset()
        # 记录走过的步数
        step_count = 0
        while True:
            # 更新可视化环境
            env.render()

            action = dqn.choose_action(state)

            state_, reward, done = env.step(action)

            step_count += 1

            dqn.store_transition(state, action, reward, state_)

            if (step > 200) and (step % 5 == 0):
                dqn.learn()
            # 将下一个 state 的值 传到下一次循环
            state = state_

            # 如果踩到 炸弹或者 宝藏 这回和结束
            if done:
                print("回合 {} 结束， 总步数： {} \n".format(episode + 1, step_count))
                break
            step += 1
    # 结束游戏 并关闭窗口
    print("游戏结束")
    env.destory()


if __name__ == '__main__':
    # 创建 环境 env  和 RL
    env = Maze()

    # 是否输出 Tensorboard 日志
    output_graph_boolean = True
    # 创建 DeepQLearning 对象
    dqn = DeepQLearning(env.n_actions,
                        env.n_features,
                        learning_rate=0.01,
                        discount_factor=0.9,
                        e_greedy=0.1,
                        replace_target_iter=200,
                        memory_size=2000,
                        output_graph=output_graph_boolean
                        )

    # 开始可视化环境
    env.after(100, update)
    env.mainloop()

    if output_graph_boolean:
        print("\t tensorboard --logdir=logs \n")
